package com.liuqi.aiagent.service.impl;

import com.alibaba.fastjson.JSONObject;
import com.liuqi.aiagent.ai.chat.ChatException;
import com.liuqi.aiagent.mapper.ChatRecordMapper;
import com.liuqi.aiagent.mapper.TokensUsageMapper;
import com.liuqi.aiagent.service.AgentChatManager;
import com.liuqi.aiagent.service.ChatService;
import com.liuqi.aiagent.service.VariableService;
import com.liuqi.aiagent.transfer.common.ChatResultMessage;
import com.liuqi.aiagent.transfer.dto.AgentChatDto;
import com.liuqi.aiagent.transfer.entity.ChatRecord;
import com.liuqi.aiagent.transfer.entity.TokensUsage;
import com.liuqi.aiagent.transfer.vo.VariableVo;
import com.liuqi.aiagent.util.IdWorker;
import com.liuqi.openai.chain.ChainStreamingHandler;
import com.liuqi.openai.chain.OpenAiChain;
import com.liuqi.openai.chain.OpenAiChainException;
import com.liuqi.openai.chain.result.ChainResult;
import com.liuqi.openai.chain.result.Content;
import com.liuqi.openai.chain.result.Function;
import com.liuqi.openai.chain.result.Think;
import com.liuqi.openai.data.message.AiMessage;
import com.liuqi.openai.data.message.ChatMessage;
import com.liuqi.openai.data.message.ToolResultMessage;
import com.liuqi.openai.model.output.TokenUsage;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

/**
 * @author liuqi
 * @date 2025/7/18
 **/
@Slf4j
@Service
@AllArgsConstructor
public class ChatServiceImpl implements ChatService {

    private final String errorResponseMessage = "{\"error\": \"系统服务异常,请联系管理员处理\"}";

    private final AgentChatManager agentChatManager;

    private final ChatRecordMapper chatRecordMapper;

    private final TokensUsageMapper tokensUsageMapper;

    private final VariableService variableService;

    /**
     * 存储用户当前的聊天
     * {
     * username: {
     * agentName: CurrentUserChat
     * }
     * }
     */
    private final ConcurrentHashMap<String, ConcurrentHashMap<String, String>> currentUserChat = new ConcurrentHashMap<>();

    /**
     * 对话异步执行器
     */
    private final Executor executor = new ThreadPoolExecutor(4, 120,
            0L, TimeUnit.MILLISECONDS,
            new LinkedBlockingQueue<Runnable>());

    @Override
    public ResponseBodyEmitter chatCompletions(AgentChatDto agentChat) throws ChatException {
        String agentName = agentChat.getAgentName();
        String username = agentChat.getUsername();

        try {
            if (currentUserChatting(agentName, username)) {
                throw new ChatException("当前用户[" + agentChat.getUsername() + "], 在Agent[" + agentChat.getAgentName() + "]有未完成的对话");
            }

            // 变量
            Map<String, String> variables = mergeVariables(agentChat.getVariables());

            // 记录用户在agent上聊天，用于标记当前正在聊天
            putCurrentUserChat(agentName, username, agentChat.getInput());


            ResponseBodyEmitter emitter = new ResponseBodyEmitter((long) (agentChat.getTimeout() * 60000));
            OpenAiChain<String> openAiChain = agentChatManager.findOpenAiChain(agentName, username);
            executor.execute(() -> {
                try {
                    // 生成对话ID
                    String uuid = IdWorker.uuidToString();

                    // 注入变量并发起AI请求
                    ChainResult chainResult = openAiChain.injectCurrentChatVariables(variables)
                            .execute(agentChat.getInput(), new ChatChainStreamingHandler(emitter));

                    log.info("对话结束,uuid[{}] 耗时[{}]", uuid, chainResult.getEndTime() - chainResult.getBeginTime());

                    // 对话消息结果处理
                    chainResultHandle(agentChat, uuid, chainResult);
                } catch (OpenAiChainException e) {
                    log.error("OpenAiChain execute error.", e);
                }
            });

            return emitter;
        } catch (Throwable t) {
            throw new ChatException("chatCompletions error", t);
        } finally {
            removeCurrentUserChat(agentName, username);
        }
    }

    /**
     * 当前用户在聊天吗
     *
     * @param agentName
     * @param username
     * @return
     */
    private boolean currentUserChatting(String agentName, String username) {
        return getCurrentUserChat(agentName, username) != null;
    }

    private String getCurrentUserChat(String agentName, String username) {
        ConcurrentHashMap<String, String> currentUserChats = currentUserChat.get(username);
        if (currentUserChats == null) {
            return null;
        }
        String currentUserChat = currentUserChats.get(agentName);
        if (currentUserChat == null) {
            return null;
        }
        return currentUserChat;
    }

    private void putCurrentUserChat(String agentName, String username, String input) {
        ConcurrentHashMap<String, String> currentUserChats = currentUserChat.get(username);
        if (currentUserChats == null) {
            synchronized (this) {
                currentUserChat.put(username, currentUserChats = new ConcurrentHashMap<>());
            }
        }
        currentUserChats.put(agentName, input);
    }

    private void removeCurrentUserChat(String agentName, String username) {
        ConcurrentHashMap<String, String> currentUserChats = currentUserChat.get(username);
        if (currentUserChats != null) {
            currentUserChats.remove(agentName);
        }
    }

    /**
     * 与系统环境变量进行合并
     *
     * @param variables
     * @return
     */
    private Map<String, String> mergeVariables(Map<String, String> variables) {
        if (variables == null) {
            variables = new HashMap<>();
        }
        List<VariableVo> variableVos = variableService.getVariables();
        for (VariableVo variableVo : variableVos) {
            if (!variables.containsKey(variableVo.getName())) {
                variables.put(variableVo.getName(), variableVo.getValue());
            }
        }
        return variables;
    }

    /**
     * 对话结果消息处理, 落库
     *
     * @param agentChat 用户发起对话请求
     * @param chatId    对话ID
     * @param result    AI响应结果
     */
    private void chainResultHandle(AgentChatDto agentChat, String chatId, ChainResult result) {
        Integer agentId = agentChatManager.findAgentId(agentChat.getAgentName(), agentChat.getUsername());
        String modelName = agentChatManager.findModelName(agentChat.getAgentName(), agentChat.getUsername());

        // 对话存储记录
        ChatRecord chatRecord = ChatRecord.builder()
                .agentId(agentId)
                .agentName(agentChat.getAgentName())
                .modelName(modelName)
                .username(agentChat.getUsername())
                .chatId(chatId)
                .beginTime(result.getBeginTime())
                .endTime(result.getEndTime())
                .input(agentChat.getInput())
                .build();

        // 对话Tokens计量
        List<TokensUsage> tokensUsages = new ArrayList<>();

        // 处理AI响应消息
        List<ChainResult.AssistantMessage> assistantMessages = result.getAssistantMessages();
        if (assistantMessages != null) {
            // 对话响应结果
            List<ChatResultMessage> chatResultMessages = new ArrayList<>();

            assistantMessages.forEach(aiMsg -> {
                ChatResultMessage.ChatResultMessageBuilder builder = ChatResultMessage.builder()
                        .startTime(aiMsg.getStartTime())
                        .endTime(aiMsg.getEndTime());

                ChatMessage chatMessage = aiMsg.getAssistantMessage();
                if (chatMessage instanceof AiMessage) {
                    // AI消息
                    AiMessage aiMessage = (AiMessage) chatMessage;

                    // 设置AI消息
                    builder.assistantMessage(
                            ChatResultMessage.AssistantMessage.builder()
                                    .id(aiMessage.id())
                                    .content(aiMessage.text())
                                    .reasoningContent(aiMessage.reasoningContent())
                                    .thinkDuration(aiMsg.getThinkDuration())
                                    .toolCallRequests(aiMessage.toolCallRequests())
                                    .build()
                    );
                } else if (chatMessage instanceof ToolResultMessage) {
                    // 工具调用结果
                    ToolResultMessage toolResult = (ToolResultMessage) chatMessage;

                    // 设置工具调用结果
                    builder.toolCallResult(
                            ChatResultMessage.ToolCallResult.builder()
                                    .name(toolResult.name())
                                    .toolCallId(toolResult.toolCallId())
                                    .result(toolResult.text())
                                    .build()
                    );
                }

                // 添加对话存储记录
                chatResultMessages.add(builder.build());

                // 处理 TokensUsage
                TokenUsage tokenUsage = aiMsg.getTokenUsage();
                if (tokenUsage != null) {
                    tokensUsages.add(
                            TokensUsage.builder()
                                    .agentName(agentChat.getAgentName())
                                    .modelName(modelName)
                                    .modelType("LLM")
                                    .username(agentChat.getUsername())
                                    .uuid(chatId)
                                    .inputToken(tokenUsage.getInputTokenCount())
                                    .outputToken(tokenUsage.getOutputTokenCount())
                                    .totalToken(tokenUsage.getTotalTokenCount())
                                    .build()
                    );
                }
            });

            // 设置响应消息结果集, json类型
            log.info("存储对话记录, 对话id[{}]", chatId);
            chatRecord.setOutput(JSONObject.toJSONString(chatResultMessages));

            if (!tokensUsages.isEmpty()) {
                log.info("存储对话使用Token计量, 对话uuid[{}]", chatId);
                tokensUsageMapper.saveBatch(tokensUsages);
            }
        }

        // 落库
        chatRecordMapper.save(chatRecord);
    }

    private class ChatChainStreamingHandler implements ChainStreamingHandler {
        final ResponseBodyEmitter emitter;

        ChatChainStreamingHandler(ResponseBodyEmitter emitter) {
            this.emitter = emitter;
        }

        @Override
        public void onContent(Content content) {
            sendData(JSONObject.toJSONString(content));
        }

        @Override
        public void onThink(Think think) {
            sendData(JSONObject.toJSONString(think));
        }

        @Override
        public void onToolCall(Function function) {
            JSONObject jsonObject = new JSONObject();
            jsonObject.put("function", function);
            sendData(jsonObject.toJSONString());
        }

        @Override
        public void onComplete() {
            sendData("[DONE]");
            emitter.complete();
        }

        @Override
        public void onError(Throwable error) {
            sendData(errorResponseMessage);
            sendData("[DONE]");
            emitter.complete();
        }

        void sendData(String data) {
            try {
                StringBuilder buf = new StringBuilder();
                buf.append("data: ")
                        .append(data)
                        .append("\n\n");
                emitter.send(buf.toString());
            } catch (Throwable t) {
                log.error("emitter 写入数据异常." + t.getMessage());
            }
        }

    }

}
